import os
import dgl
import itertools
import numpy as np
import pandas as pd

import torch
from torch.utils.data import Dataset, DataLoader

from Config import Config
from data.data_parser import GraphParser
from data.pos_encode import positional_encoding
from data.preprocess import del_errors

"""
将数据集封装为torch中的Dataset类型
方便使用DataLoader批量读取数据进行训练
"""


class mRNADataset(Dataset):
    def __init__(self, data_file, config=None, mode='train'):
        self.data_file = data_file
        self.config = config
        self.mode = mode

        df = pd.read_json(self.data_file, lines=True)
        # 直接将噪音项删除
        df = del_errors(df, max_error=5)
        if self.mode == 'train':
            # 若是训练模式，则打乱数据集
            df = df.sample(frac=1)
            df = df.reset_index(drop=True)

        self.parser = GraphParser(self.config, mode='train')

        # preprocess data
        data = []
        graph_list = []
        node_labels = []
        for i in range(len(df)):
            sample = df.loc[i]
            data.append(self.parser.parse(sample))

            # 向graph_list中添加图
            gdata = self.parser.parse(sample)
            # 构造图
            graph = dgl.DGLGraph()
            num_nodes = len(gdata['nfeat'])
            graph.add_nodes(num_nodes)
            graph.add_edges(gdata['edges'][:, 0], gdata['edges'][:, 1])
            # 给每个图的特征赋值
            graph.ndata['nfeat'] = torch.tensor(gdata['nfeat'])
            graph.ndata['loop_type'] = torch.tensor(gdata['loop_type'])
            graph.ndata['mask'] = torch.tensor(gdata['mask'])
            graph.edata['efeat'] = torch.tensor(gdata['efeat'])

            node_labels.append(gdata['labels'])

            # 将配置好的图放入graph_list
            graph_list.append(graph)

        self.data = data
        self.graph_list = graph_list
        self.node_labels = node_labels

    def __getitem__(self, idx):
        return self.graph_list[idx], self.node_labels[idx]

    def __len__(self):
        return len(self.graph_list)

    def _add_positional_encodings(self, pos_enc_dim):
        # 添加位置节点
        self.graph_list = [positional_encoding(g, pos_enc_dim) for g in self.graph_list]


class collate_fn(object):
    def __init__(self, config=None, mode='train'):
        self.config = config
        self.mode = mode

    def __call__(self, batch_data):
        # The input samples is a list of pairs (graph, label).
        # 输入是一个mini_batch的list, (graph, label)
        graphs, labels = map(list, zip(*batch_data))
        # Edge classification labels need to be flattened to 1D lists
        # 真实标签需要进行展平
        labels = torch.Tensor(np.array(list(itertools.chain(*labels))))

        batched_graph = dgl.batch(graphs)

        return batched_graph, labels


def gpu_setup(use_gpu, gpu_id):
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    if torch.cuda.is_available() and use_gpu:
        print('cuda available with GPU:', torch.cuda.get_device_name(0))
        device = torch.device("cuda")
    else:
        print('cuda not available')
        device = torch.device("cpu")
    return device


if __name__ == '__main__':
    model_name = 'GCN'
    data_file = "../dataset/valid.json"

    """GPU设置"""
    device = gpu_setup(False, 0)
    config = Config(model_name, device)
    data_set = mRNADataset(data_file, config, mode='train')

    # data_set._add_positional_encodings(pos_enc_dim=64)

    data_loader = DataLoader(data_set, batch_size=config.batch_size, shuffle=True, collate_fn=collate_fn())

    for data_iter in data_loader:
        batch_graph, node_labels = data_iter
        print(batch_graph)
